#include <stdio.h>
#include <math.h>

#include "slm.h"

bool CSlm::Load(const char* fname) {
  bool suc = false;
  FILE *fp = fopen(fname, "rb");
  if (fp != NULL && fread(&N, sizeof(N), 1, fp) == 1 &&
      fread(&bUseLogPr, sizeof(bUseLogPr), 1, fp) == 1) {
    sz = new int[N+1];
    level = new void* [N+1];
    fread(sz, sizeof(int), N+1, fp);
    for (int lvl = 0; lvl <= N; ++lvl) {
      if (lvl < N) {            // Node
        level[lvl] = new TNode[sz[lvl]];
        fread(level[lvl], sizeof(TNode), sz[lvl], fp);
      } else {                  // Leaf node
        level[lvl] = new TLeaf[sz[lvl]];
        fread(level[lvl], sizeof(TLeaf), sz[lvl], fp);
      }
    }
    suc = true;
  }
  fclose(fp);
  return suc;
}

void CSlm::Free() {
  delete [] sz;
  for (int lvl = 0; lvl <= N; ++lvl) {
    if (lvl == N)
      delete [] ((TLeaf *)level[lvl]);
    else
      delete [] ((TNode *)level[lvl]);
  }
  level = NULL;
  sz = NULL;
  bUseLogPr = 0;
  N = 0;
}

double CSlm::getNegLogPr(int n, TWORDID *hw) {
  double val = (bUseLogPr)?(getPrAsLog(n, hw)):(getPrDirect(n, hw));
  return (bUseLogPr)?(val):(-log(val));
}

double CSlm::getPr(int n, TWORDID *hw) {
  double val = (bUseLogPr)?(getPrAsLog(n, hw)):(getPrDirect(n, hw));
  return (bUseLogPr)?(exp(-val)):(val);
}

double CSlm::getPrAsLog(int n, TWORDID *hw) {
  void *pstate = ((TNode *)level[0]);
  double bow = 0.0;

  if (n > N) {
    hw += (N-n);
    n = N;
  }
  for (int lvl = 0; lvl < n && pstate != NULL; ++lvl) {
    int h = ((TNode *)pstate)->child;
    int t = ((TNode *)pstate+1)->child;

    if (lvl == n -1)
      bow = ((TNode *)pstate)->bow;
    if (lvl == N-1) {
      TLeaf *p = (TLeaf *)level[lvl+1];
      pstate = (void *)binary_find_id(p+h, p+t, hw[lvl]);
    } else {
      TNode *p = (TNode *)level[lvl+1];
      pstate = (void)binary_find_id(p+h, p+t, hw[lvl]);
    }
  }
  if (pstate == NULL) {
    return bow + getPrAsLog(n-1, hw+1);
  } else {
    return ((TLeaf *)pstate)->pr; // as we derive TNode from TLeaf
  }
}
    
double CSlm::getPrDirect(int n, TWORDID *hw) {
  void *pstate = ((TNode *)level[0]);
  double bow = 1.0;
  if (n > N) {
    hw += (N - n);
    n = N;
  }

  for (int lvl = 0; lvl < n && pstate != NULL; ++lvl) {
    int h = ((TNode *)pstate)->child;
    int t = ((TNode *)pstate+1)->child;
    if (lvl == n-1)
      bow = ((TNode *)pstate)->bow;
    if (lvl == N-1) {
      TLeaf *p = (TLeaf *)level[lvl+1];
      pstate = (void *)binary_find_id(p+h, p+t, hw[lvl]);
    } else {
      TNode *p = (TNode *)level[lvl+1];
      pstate = (void *)binary_find_id(p+h, p+t, hw[lvl]);
    }
  }
  if (pstate == NULL) {
    return bow*getPrDirect(n-1, hw+1);
  } else {
    return ((TLeaf *)pstate)->pr; // as we derive TNode from TLeaf
  }
}
